#!/usr/bin/env python

import asyncio
import logging
import websockets
import json # 现在会用到 JSON

# --- 配置 ---
HOST = '0.0.0.0'
PORT = 8765
LOG_LEVEL = logging.INFO

# --- 全局变量 ---
# 使用字典存储客户端信息 {websocket: {"username": str, "addr": tuple}}
client_info = {}
# 集合仍然用于快速迭代所有连接
connected_clients = set()

# --- 日志设置 ---
logging.basicConfig(level=LOG_LEVEL, format='%(asctime)s - %(levelname)s - %(message)s')

async def broadcast(message_data: dict):
    """
    将结构化消息 (字典) 广播给所有已认证 (有用户名) 的客户端。
    """
    if not client_info: # 检查是否有已设置用户名的客户端
        logging.info("No authenticated clients to broadcast to.")
        return

    message_str = json.dumps(message_data) # 将字典转换为 JSON 字符串
    clients_to_send = connected_clients.copy() # 迭代副本
    logging.info(f"Broadcasting {message_str} to {len(clients_to_send)} client(s).")

    disconnected_clients = set()
    for client in clients_to_send:
        # 只向已经设置了用户名的客户端广播（避免给刚连接还没发用户名的客户端发消息）
        if client in client_info:
            try:
                await client.send(message_str)
            except websockets.exceptions.ConnectionClosedOK:
                logging.info(f"Client {client_info.get(client, {}).get('addr', 'Unknown')} disconnected cleanly during broadcast.")
                disconnected_clients.add(client)
            except websockets.exceptions.ConnectionClosedError:
                logging.warning(f"Client {client_info.get(client, {}).get('addr', 'Unknown')} connection closed with error during broadcast.")
                disconnected_clients.add(client)
            except Exception as e:
                logging.error(f"Error sending message to {client_info.get(client, {}).get('addr', 'Unknown')}: {e}")
                disconnected_clients.add(client)

    # 清理广播时发现断开的客户端
    for client in disconnected_clients:
        if client in connected_clients:
            connected_clients.remove(client)
        if client in client_info:
            logging.info(f"Removed disconnected client {client_info[client]['addr']} (Username: {client_info[client]['username']}) after broadcast failure.")
            del client_info[client] # 同时从 client_info 中移除


async def handle_client(websocket: websockets.WebSocketServerProtocol):
    """
    处理单个客户端连接。
    """
    client_addr = websocket.remote_address
    logging.info(f"Client connected: {client_addr}. Waiting for username.")

    # 1. 将新连接加入 connected_clients 集合，但还不加入 client_info
    connected_clients.add(websocket)
    username = None # 初始用户名未知

    try:
        # 2. 等待并处理第一条消息（期望是用户名设置）
        first_message = await websocket.recv()
        logging.info(f"Received first message from {client_addr}: {first_message}")
        try:
            data = json.loads(first_message)
            if data.get("type") == "set_username" and "username" in data:
                username = data["username"].strip()
                if not username:
                    await websocket.close(code=1008, reason="Username cannot be empty.")
                    logging.warning(f"Client {client_addr} provided empty username. Connection closed.")
                    return # 结束此 handler

                # 检查用户名是否已存在（简单示例，允许重复）
                # if any(info["username"] == username for info in client_info.values()):
                #     await websocket.close(code=1008, reason="Username already taken.")
                #     logging.warning(f"Client {client_addr} tried existing username '{username}'. Connection closed.")
                #     return

                # 用户名有效，存储信息并广播加入消息
                client_info[websocket] = {"username": username, "addr": client_addr}
                logging.info(f"Client {client_addr} set username to '{username}'. Total authenticated clients: {len(client_info)}")

                join_message = {
                    "type": "system",
                    "text": f"用户 '{username}' 加入了聊天室。"
                }
                await broadcast(join_message) # 广播加入消息
            else:
                # 第一条消息格式不对
                await websocket.close(code=1008, reason="Invalid initial message format. Expected username setup.")
                logging.warning(f"Client {client_addr} sent invalid first message. Connection closed.")
                return
        except json.JSONDecodeError:
            await websocket.close(code=1008, reason="Invalid JSON format in initial message.")
            logging.warning(f"Client {client_addr} sent non-JSON first message. Connection closed.")
            return
        except Exception as e:
            await websocket.close(code=1011, reason=f"Error processing username: {e}")
            logging.error(f"Error processing username for {client_addr}: {e}")
            return

        # 3. 用户名设置成功后，循环接收聊天消息
        async for message in websocket:
            logging.info(f"Received chat message from '{username}' ({client_addr}): {message}")

            # 4. 构建聊天消息 JSON 并广播
            chat_message = {
                "type": "chat",
                "sender": username, # 使用设置好的用户名
                "text": message
            }
            await broadcast(chat_message)

    except websockets.exceptions.ConnectionClosedOK:
        logging.info(f"Client '{username}' ({client_addr}) disconnected cleanly.")
    except websockets.exceptions.ConnectionClosedError:
        logging.warning(f"Client '{username}' ({client_addr}) connection closed with error.")
    except Exception as e:
        logging.error(f"An unexpected error occurred with client '{username}' ({client_addr}): {e}")
    finally:
        # 5. 清理连接信息并广播离开消息
        if websocket in connected_clients:
             connected_clients.remove(websocket)

        if websocket in client_info:
             removed_username = client_info[websocket]["username"]
             logging.info(f"Client '{removed_username}' ({client_addr}) removed from client_info.")
             del client_info[websocket] # 从 client_info 中移除
             logging.info(f"Total authenticated clients: {len(client_info)}")

             # 广播离开消息
             leave_message = {
                 "type": "system",
                 "text": f"用户 '{removed_username}' 离开了聊天室。"
             }
             # 确保在 finally 中也能安全地广播
             try:
                 # 使用 asyncio.create_task 以免阻塞 finally
                 asyncio.create_task(broadcast(leave_message))
             except Exception as broadcast_err:
                 logging.error(f"Error scheduling broadcast for leave message of '{removed_username}': {broadcast_err}")
        else:
            # 如果客户端在设置用户名之前就断开了
            logging.info(f"Unauthenticated client {client_addr} disconnected.")


async def main():
    async with websockets.serve(handle_client, HOST, PORT):
        logging.info(f"WebSocket server started on ws://{HOST}:{PORT}")
        await asyncio.Future()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logging.info("Server stopped by user (Ctrl+C).")
    except Exception as e:
        logging.error(f"Server encountered a fatal error: {e}")